package gbase

import (
	"cmp"
	"fmt"
)

// ==================== 泛型堆接口 ====================

// Heap 泛型堆接口
type Heap[T any] interface {
	Push(x T)
	Pop() T
	Peek() T
	Len() int
	IsEmpty() bool
	Clear()
	ToSlice() []T
}

// ==================== 最小堆实现 ====================

// MinHeap 最小堆
type MinHeap[T cmp.Ordered] struct {
	data []T
}

// NewMinHeap 创建最小堆
func NewMinHeap[T cmp.Ordered](capacity ...int) *MinHeap[T] {
	cap := 0
	if len(capacity) > 0 {
		cap = capacity[0]
	}
	return &MinHeap[T]{
		data: make([]T, 0, cap),
	}
}

// Push 添加元素
func (h *MinHeap[T]) Push(x T) {
	h.data = append(h.data, x)
	h.up(len(h.data) - 1)
}

// Pop 移除并返回最小元素
func (h *MinHeap[T]) Pop() T {
	if h.IsEmpty() {
		var zero T
		return zero
	}

	n := len(h.data) - 1
	h.swap(0, n)
	h.down(0, n)

	val := h.data[n]
	h.data = h.data[:n]
	return val
}

// Peek 查看最小元素（不移除）
func (h *MinHeap[T]) Peek() T {
	if h.IsEmpty() {
		var zero T
		return zero
	}
	return h.data[0]
}

// Len 返回堆大小
func (h *MinHeap[T]) Len() int {
	return len(h.data)
}

// IsEmpty 检查是否为空
func (h *MinHeap[T]) IsEmpty() bool {
	return len(h.data) == 0
}

// Clear 清空堆
func (h *MinHeap[T]) Clear() {
	h.data = h.data[:0]
}

// ToSlice 转换为切片
func (h *MinHeap[T]) ToSlice() []T {
	result := make([]T, len(h.data))
	copy(result, h.data)
	return result
}

// up 向上调整
func (h *MinHeap[T]) up(j int) {
	for {
		parent := (j - 1) / 2
		if j == parent || h.data[parent] <= h.data[j] {
			break
		}
		h.swap(parent, j)
		j = parent
	}
}

// down 向下调整
func (h *MinHeap[T]) down(i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 {
			break
		}
		j := j1
		if j2 := j1 + 1; j2 < n && h.data[j2] < h.data[j1] {
			j = j2
		}
		if h.data[j] >= h.data[i] {
			break
		}
		h.swap(i, j)
		i = j
	}
	return i > i0
}

// swap 交换元素
func (h *MinHeap[T]) swap(i, j int) {
	h.data[i], h.data[j] = h.data[j], h.data[i]
}

// ==================== 最大堆实现 ====================

// MaxHeap 最大堆
type MaxHeap[T cmp.Ordered] struct {
	data []T
}

// NewMaxHeap 创建最大堆
func NewMaxHeap[T cmp.Ordered](capacity ...int) *MaxHeap[T] {
	cap := 0
	if len(capacity) > 0 {
		cap = capacity[0]
	}
	return &MaxHeap[T]{
		data: make([]T, 0, cap),
	}
}

// Push 添加元素
func (h *MaxHeap[T]) Push(x T) {
	h.data = append(h.data, x)
	h.up(len(h.data) - 1)
}

// Pop 移除并返回最大元素
func (h *MaxHeap[T]) Pop() T {
	if h.IsEmpty() {
		var zero T
		return zero
	}

	n := len(h.data) - 1
	h.swap(0, n)
	h.down(0, n)

	val := h.data[n]
	h.data = h.data[:n]
	return val
}

// Peek 查看最大元素（不移除）
func (h *MaxHeap[T]) Peek() T {
	if h.IsEmpty() {
		var zero T
		return zero
	}
	return h.data[0]
}

// Len 返回堆大小
func (h *MaxHeap[T]) Len() int {
	return len(h.data)
}

// IsEmpty 检查是否为空
func (h *MaxHeap[T]) IsEmpty() bool {
	return len(h.data) == 0
}

// Clear 清空堆
func (h *MaxHeap[T]) Clear() {
	h.data = h.data[:0]
}

// ToSlice 转换为切片
func (h *MaxHeap[T]) ToSlice() []T {
	result := make([]T, len(h.data))
	copy(result, h.data)
	return result
}

// up 向上调整
func (h *MaxHeap[T]) up(j int) {
	for {
		parent := (j - 1) / 2
		if j == parent || h.data[parent] >= h.data[j] {
			break
		}
		h.swap(parent, j)
		j = parent
	}
}

// down 向下调整
func (h *MaxHeap[T]) down(i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 {
			break
		}
		j := j1
		if j2 := j1 + 1; j2 < n && h.data[j2] > h.data[j1] {
			j = j2
		}
		if h.data[j] <= h.data[i] {
			break
		}
		h.swap(i, j)
		i = j
	}
	return i > i0
}

// swap 交换元素
func (h *MaxHeap[T]) swap(i, j int) {
	h.data[i], h.data[j] = h.data[j], h.data[i]
}

// ==================== 自定义比较器堆 ====================

// ComparatorHeap 支持自定义比较器的堆
type ComparatorHeap[T any] struct {
	data       []T
	comparator func(a, b T) int // 返回 -1, 0, 1 表示 a < b, a == b, a > b
	isMin      bool             // true 为最小堆，false 为最大堆
}

// NewComparatorHeap 创建自定义比较器堆
func NewComparatorHeap[T any](comparator func(a, b T) int, isMin bool, capacity ...int) *ComparatorHeap[T] {
	cap := 0
	if len(capacity) > 0 {
		cap = capacity[0]
	}
	return &ComparatorHeap[T]{
		data:       make([]T, 0, cap),
		comparator: comparator,
		isMin:      isMin,
	}
}

// Push 添加元素
func (h *ComparatorHeap[T]) Push(x T) {
	h.data = append(h.data, x)
	h.up(len(h.data) - 1)
}

// Pop 移除并返回堆顶元素
func (h *ComparatorHeap[T]) Pop() T {
	if h.IsEmpty() {
		var zero T
		return zero
	}

	n := len(h.data) - 1
	h.swap(0, n)
	h.down(0, n)

	val := h.data[n]
	h.data = h.data[:n]
	return val
}

// Peek 查看堆顶元素（不移除）
func (h *ComparatorHeap[T]) Peek() T {
	if h.IsEmpty() {
		var zero T
		return zero
	}
	return h.data[0]
}

// Len 返回堆大小
func (h *ComparatorHeap[T]) Len() int {
	return len(h.data)
}

// IsEmpty 检查是否为空
func (h *ComparatorHeap[T]) IsEmpty() bool {
	return len(h.data) == 0
}

// Clear 清空堆
func (h *ComparatorHeap[T]) Clear() {
	h.data = h.data[:0]
}

// ToSlice 转换为切片
func (h *ComparatorHeap[T]) ToSlice() []T {
	result := make([]T, len(h.data))
	copy(result, h.data)
	return result
}

// up 向上调整
func (h *ComparatorHeap[T]) up(j int) {
	for {
		parent := (j - 1) / 2
		if j == parent {
			break
		}

		shouldSwap := false
		if h.isMin {
			shouldSwap = h.comparator(h.data[parent], h.data[j]) > 0
		} else {
			shouldSwap = h.comparator(h.data[parent], h.data[j]) < 0
		}

		if !shouldSwap {
			break
		}

		h.swap(parent, j)
		j = parent
	}
}

// down 向下调整
func (h *ComparatorHeap[T]) down(i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 {
			break
		}
		j := j1
		if j2 := j1 + 1; j2 < n {
			shouldSwap := false
			if h.isMin {
				shouldSwap = h.comparator(h.data[j2], h.data[j1]) < 0
			} else {
				shouldSwap = h.comparator(h.data[j2], h.data[j1]) > 0
			}
			if shouldSwap {
				j = j2
			}
		}

		shouldSwap := false
		if h.isMin {
			shouldSwap = h.comparator(h.data[j], h.data[i]) < 0
		} else {
			shouldSwap = h.comparator(h.data[j], h.data[i]) > 0
		}

		if !shouldSwap {
			break
		}

		h.swap(i, j)
		i = j
	}
	return i > i0
}

// swap 交换元素
func (h *ComparatorHeap[T]) swap(i, j int) {
	h.data[i], h.data[j] = h.data[j], h.data[i]
}

// ==================== 堆排序 ====================

// HeapSort 堆排序（原地排序）
func HeapSort[T cmp.Ordered](data []T, ascending bool) {
	if len(data) <= 1 {
		return
	}

	n := len(data)

	// 构建堆
	for i := n/2 - 1; i >= 0; i-- {
		if ascending {
			heapDownMax(data, i, n) // 升序用最大堆
		} else {
			heapDownMin(data, i, n) // 降序用最小堆
		}
	}

	// 排序：将堆顶元素移到末尾，然后重新调整堆
	for i := n - 1; i > 0; i-- {
		data[0], data[i] = data[i], data[0]
		if ascending {
			heapDownMax(data, 0, i) // 升序用最大堆
		} else {
			heapDownMin(data, 0, i) // 降序用最小堆
		}
	}
}

// heapDownMin 最小堆向下调整
func heapDownMin[T cmp.Ordered](data []T, i0, n int) {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 {
			break
		}
		j := j1
		if j2 := j1 + 1; j2 < n && data[j2] < data[j1] {
			j = j2
		}
		if data[j] >= data[i] {
			break
		}
		data[i], data[j] = data[j], data[i]
		i = j
	}
}

// heapDownMax 最大堆向下调整
func heapDownMax[T cmp.Ordered](data []T, i0, n int) {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 {
			break
		}
		j := j1
		if j2 := j1 + 1; j2 < n && data[j2] > data[j1] {
			j = j2
		}
		if data[j] <= data[i] {
			break
		}
		data[i], data[j] = data[j], data[i]
		i = j
	}
}

// ==================== 工具函数 ====================

// BuildHeap 从切片构建堆
func BuildMinHeap[T cmp.Ordered](data []T) *MinHeap[T] {
	h := NewMinHeap[T](len(data))
	for _, v := range data {
		h.Push(v)
	}
	return h
}

// BuildMaxHeap 从切片构建最大堆
func BuildMaxHeap[T cmp.Ordered](data []T) *MaxHeap[T] {
	h := NewMaxHeap[T](len(data))
	for _, v := range data {
		h.Push(v)
	}
	return h
}

// TopK 获取前 K 个最大/最小元素
func TopK[T cmp.Ordered](data []T, k int, largest bool) []T {
	if k <= 0 || len(data) == 0 {
		return nil
	}

	if k >= len(data) {
		result := make([]T, len(data))
		copy(result, data)
		return result
	}

	var h Heap[T]
	if largest {
		h = NewMinHeap[T](k)
	} else {
		h = NewMaxHeap[T](k)
	}

	for _, v := range data {
		if h.Len() < k {
			h.Push(v)
		} else {
			if largest {
				if v > h.Peek() {
					h.Pop()
					h.Push(v)
				}
			} else {
				if v < h.Peek() {
					h.Pop()
					h.Push(v)
				}
			}
		}
	}

	return h.ToSlice()
}

// ==================== 调试和可视化 ====================

// PrintHeap 打印堆结构（用于调试）
func PrintHeap[T any](h Heap[T], title string) {
	fmt.Printf("=== %s ===\n", title)
	fmt.Printf("Size: %d\n", h.Len())
	if h.IsEmpty() {
		fmt.Println("Empty heap")
		return
	}

	data := h.ToSlice()
	for i, v := range data {
		fmt.Printf("[%d] %v\n", i, v)
	}
	fmt.Println()
}
